from argparse import ArgumentParser
import itertools
import os

import arff

from autosklearn.constants import *
from autosklearn.metrics import CLASSIFICATION_METRICS, REGRESSION_METRICS

if __name__ == "__main__":
    parser = ArgumentParser()
    parser.add_argument("--working-directory", type=str, required=True)
    parser.add_argument("--task-type", required=True,
                        choices=['classification', 'regression'])
    parser.add_argument("--scenario_id", type=str, default='auto-sklearn')
    parser.add_argument("--algorithm_cutoff_time", type=int, default=1800)
    parser.add_argument("--algorithm_cutoff_memory", type=int, default=3072)

    args = parser.parse_args()
    working_directory = args.working_directory
    task_type = args.task_type

    output_dir = os.path.join(working_directory, 'metadata')
    results_dir = os.path.join(working_directory, 'configuration_results')
    metafeatures_dir = os.path.join(working_directory, 'metafeatures')

    scenario_id = args.scenario_id
    algorithm_cutoff_time = args.algorithm_cutoff_time
    algorithm_cutoff_memory = args.algorithm_cutoff_memory

    # Create the output directory if necessary
    try:
        os.makedirs(output_dir)
    except (OSError, IOError):
        pass

    if task_type == 'classification':
        metadata_sets = itertools.product(
            [0, 1], [BINARY_CLASSIFICATION, MULTICLASS_CLASSIFICATION],
            CLASSIFICATION_METRICS)
        input_directory = os.path.join(working_directory, 'configuration',
                                       'classification')
    elif task_type == 'regression':
        metadata_sets = itertools.product(
            [0, 1], [REGRESSION], REGRESSION_METRICS)
        input_directory = os.path.join(working_directory, 'configuration',
                                       'regression')
    else:
        raise ValueError(task_type)

    for sparse, task, metric in metadata_sets:
        print(TASK_TYPES_TO_STRING[task], metric, sparse)

        dir_name = '%s_%s_%s' % (metric, TASK_TYPES_TO_STRING[task],
                                 'sparse' if sparse else 'dense')
        output_dir_ = os.path.join(output_dir, dir_name)
        results_dir_ = os.path.join(results_dir, dir_name)

        if not os.path.exists(results_dir_):
            print("Results directory %s does not exist!" % results_dir_)
            continue

        try:
            os.makedirs(output_dir_)
        except Exception:
            pass

        # Create a readme.txt
        with open(os.path.join(output_dir_, "readme.txt"), "w") as fh:
            pass

        # Create description.txt
        with open(os.path.join(metafeatures_dir,
                               "description.features.txt")) as fh:
            description_metafeatures = fh.read()

        with open(os.path.join(results_dir_,
                               "description.results.txt")) as fh:
            description_results = fh.read()

        description = [description_metafeatures, description_results]
        description.append("scenario_id: %s" % scenario_id)
        description.append("maximize: false")
        description.append(
            "algorithm_cutoff_time: %d" % algorithm_cutoff_time)
        description.append(
            "algorithm_cutoff_memory: %d" % algorithm_cutoff_memory)

        with open(os.path.join(output_dir_, "description.txt"), "w") as fh:
            for line in description:
                fh.write(line)
                fh.write("\n")

        # Copy feature values and add instance id
        with open(os.path.join(metafeatures_dir,
                               "feature_values.arff")) as fh:
            feature_values = arff.load(fh)

        feature_values['relation'] = scenario_id + "_" + feature_values[
            'relation']

        with open(os.path.join(output_dir_, "feature_values.arff"),
                  "w") as fh:
            arff.dump(feature_values, fh)

        # Copy feature runstatus and add instance id
        with open(os.path.join(metafeatures_dir,
                               "feature_runstatus.arff")) as fh:
            feature_runstatus = arff.load(fh)

        feature_runstatus['relation'] = scenario_id + "_" + \
                                        feature_runstatus['relation']

        with open(os.path.join(output_dir_, "feature_runstatus.arff"), "w") \
                as fh:
            arff.dump(feature_runstatus, fh)

        # Copy feature runstatus and add instance id
        with open(
                os.path.join(metafeatures_dir, "feature_costs.arff")) as fh:
            feature_costs = arff.load(fh)

        feature_costs['relation'] = scenario_id + "_" + feature_costs[
            'relation']
        for i in range(len(feature_costs['data'])):
            for j in range(2, len(feature_costs['data'][i])):
                feature_costs['data'][i][j] = \
                    round(feature_costs['data'][i][j], 5)

        with open(os.path.join(output_dir_, "feature_costs.arff"), "w") \
                as fh:
            arff.dump(feature_costs, fh)

        # Copy algorithm runs and add instance id
        with open(os.path.join(results_dir_, "algorithm_runs.arff")) as fh:
            algorithm_runs = arff.load(fh)

        algorithm_runs['relation'] = scenario_id + "_" + algorithm_runs[
            'relation']

        with open(os.path.join(output_dir_, "algorithm_runs.arff"), "w") \
                as fh:
            arff.dump(algorithm_runs, fh)

        # Copy configurations file
        with open(os.path.join(results_dir_, "configurations.csv")) as fh:
            algorithm_runs = fh.read()
        with open(os.path.join(output_dir_, "configurations.csv"), "w") \
                as fh:
            fh.write(algorithm_runs)
